# -*- coding: utf-8 -*-
# @Function: 数据库管理器
# @Description: 负责数据库连接、表操作和数据查询，提供数据库层面的CRUD操作
# @Usage: 被importer.py、query_tester.py和llm_query_tester.py导入使用，处理数据库操作
# @Dependencies: Z_config.py

import psycopg2
from psycopg2 import sql
from psycopg2.extras import Json
import json
from typing import List, Tuple
from .config import DB_CONFIG

class DBManager:
    def __init__(self, db_config=None):
        """初始化数据库管理器

        Args:
            db_config (dict, optional): 数据库配置. Defaults to None.
        """
        self.db_config = db_config or DB_CONFIG
        self.conn = None

    def connect(self):
        """连接数据库"""
        try:
            if self.conn is None or self.conn.closed != 0:
                # 强制设置环境编码为UTF-8，解决Windows编码问题
                import os
                os.environ['PGCLIENTENCODING'] = 'UTF8'
                
                self.conn = psycopg2.connect(**self.db_config)
                
                # 连接后再次确认编码
                with self.conn.cursor() as cursor:
                    cursor.execute("SET CLIENT_ENCODING TO 'UTF8';")
                self.conn.commit()
                
                print("数据库连接成功")
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            print("数据库连接已关闭")

    def drop_table(self):
        """删除QA知识库表"""
        try:
            with self.conn.cursor() as cursor:
                cursor.execute("DROP TABLE IF EXISTS qa_knowledge;")
                self.conn.commit()
                print("QA知识库表已删除")
        except Exception as e:
            print(f"删除表失败: {e}")
            raise

    def create_table(self):
        """创建QA知识库表和表格存储表"""
        self.drop_table()  # 先删除旧表
        # QA知识库表
        create_qa_table_sql = """
        CREATE TABLE qa_knowledge (
            question TEXT PRIMARY KEY,
            question_tokens JSONB,
            answer JSONB,
            context TEXT,
            category_id VARCHAR(50),
            created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
        );
        """

        # 创建索引以提高查询性能
        create_index_sql = """
        CREATE INDEX IF NOT EXISTS idx_qa_question_tokens_gin 
        ON qa_knowledge USING GIN (question_tokens);

        CREATE INDEX IF NOT EXISTS idx_qa_tokens_array 
        ON qa_knowledge USING GIN ((question_tokens->'tokens'));

        CREATE INDEX IF NOT EXISTS idx_qa_category_id 
        ON qa_knowledge (category_id);
        """

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(create_qa_table_sql)
                cursor.execute(create_index_sql)
                self.conn.commit()
                print("QA知识库表检查/创建成功")
                return True
        except Exception as e:
            print(f"创建表失败: {e}")
            raise

    def insert_qa_data(self, qa_pairs: List[Tuple[str, str, dict, str]], category_id: str):
        """插入QA数据到数据库

        Args:
            qa_pairs (List[Tuple[str, str, dict, str]]): 问答对列表，每项包含(问题, 答案, 分词结果, 上下文)
            category_id (str): 问题所属的类别ID
        """
        insert_sql = """
        INSERT INTO qa_knowledge (question, question_tokens, answer, context, category_id)
        VALUES (%s, %s, %s, %s, %s)
        ON CONFLICT (question) DO UPDATE SET
            question_tokens = EXCLUDED.question_tokens,
            answer = EXCLUDED.answer,
            context = EXCLUDED.context,
            category_id = EXCLUDED.category_id,
            updated_at = CURRENT_TIMESTAMP;
        """

        try:
            with self.conn.cursor() as cursor:
                for qa_data in qa_pairs:
                    question, answer, token_data, context = qa_data
                    
                    # 构建答案JSON
                    answer_json = {"answer": answer}

                    # 插入数据
                    cursor.execute(insert_sql, (
                        question,
                        Json(token_data),
                        Json(answer_json),
                        context,
                        category_id
                    ))

                self.conn.commit()
                print(f"成功插入或更新 {len(qa_pairs)} 条QA数据")
                return True
        except Exception as e:
            print(f"插入数据失败: {e}")
            self.conn.rollback()
            raise

    def query_qa(self, tokens: List[str], limit: int = 10, category_id: str = None):
        """根据分词结果查询QA数据
    
        Args:
            tokens (List[str]): 分词结果
            limit (int, optional): 返回结果数量限制. Defaults to 5.
            category_id (str, optional): 水库类别ID. Defaults to None.
    
        Returns:
            list: 查询结果
        """
        # 基础查询SQL
        query_sql = """
        SELECT question, question_tokens, answer, category_id, context,
               (
                   SELECT COUNT(*)
                   FROM jsonb_array_elements_text(question_tokens->'tokens') AS token
                   WHERE token = ANY(%s)
               ) as match_count
        FROM qa_knowledge
        WHERE question_tokens->'tokens' ?| %s
        """
        
        # 如果指定了category_id，添加过滤条件
        params = [tokens, tokens]
        if category_id:
            query_sql += " AND category_id = %s"
            params.append(category_id)
            print(f"按category_id过滤: {category_id}")
        
        # 添加排序和限制
        query_sql += """
        ORDER BY match_count DESC
        LIMIT %s
        """
        params.append(limit)
    
        try:
            with self.conn.cursor() as cursor:
                cursor.execute(query_sql, tuple(params))
                results = cursor.fetchall()
                return results
        except Exception as e:
            print(f"查询失败: {e}")
            return []

    def get_table_images(self, table_name: str):
        """查询 table_images 表并返回 images 字段解析后的列表

        Args:
            table_name (str): 表名（在 table_images.table_name 列中匹配）

        Returns:
            list: images 列解析后的列表（每项通常为 dict），若未找到或发生错误则返回空列表
        """
        try:
            with self.conn.cursor() as cursor:
                cursor.execute("SELECT images FROM table_images WHERE table_name = %s", (table_name,))
                row = cursor.fetchone()
                if not row:
                    return []

                images_data = row[0]
                # images_data 可能为 JSON 字符串或已解析的 Python 对象
                if isinstance(images_data, str):
                    try:
                        images = json.loads(images_data)
                    except Exception:
                        images = []
                else:
                    images = images_data or []

                # 确保返回列表类型
                if not isinstance(images, list):
                    return []

                return images
        except Exception as e:
            print(f"get_table_images 查询失败: {e}")
            return []